"""
Partial Least Squares Regression
"""

from sklearn.cross_decomposition import PLSRegression
from sklearn.model_selection import GridSearchCV
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import r2_score
import numpy as np
import matplotlib.pyplot as plt

plt.style.use(['science', 'grid', 'muted'])

# 导入数据集
data = np.loadtxt('data/linnerud.txt', delimiter='\t')
X, y = data[:, :3], data[:, -3:]
# 标准化指标变量和变量
X = StandardScaler().fit_transform(X)
y = StandardScaler().fit_transform(y)

# 回归模型
model_setup = PLSRegression(scale=True)
param_grid = {'n_components': range(1, 4)}

# 自动调参
gsearch = GridSearchCV(model_setup, param_grid)

# 训练模型
model = gsearch.fit(X, y)

# 预测
pred = model.predict(X)

# 绘图
plt.figure(figsize=(12, 8), dpi=300)
plt.suptitle('PLSRegression($R^2$: {:.2})'.format(r2_score(y, pred)))
style = ['r+', 'b*', 'go']
for i in range(pred.shape[1]):
    ax = plt.subplot(1, 3, i+1)
    plt.title('y%d'%(i+1))
    plt.xlabel('Prediction')
    plt.ylabel('Groud Truth')
    plt.plot(pred[:, i], y[:, i], style[i])
    plt.plot(plt.xlim(), plt.ylim(), '-')

plt.savefig('第11章：偏最小二乘回归分析/Ground Truth - Prediction')
plt.show()
